# uncompyle6 version 3.6.1
# Python bytecode 2.7 (62211)
# Decompiled from: Python 3.6.5 |Anaconda, Inc.| (default, Mar 29 2018, 13:32:41) [MSC v.1900 64 bit (AMD64)]
# Embedded file name: abcft_algorithm_deploy_renshou/utils.py
# Compiled at: 2020-04-27 21:44:40
import io, cv2, os, numpy as np, math, shutil
from PIL import Image, ImageFile
from skimage.measure import label, regionprops
from six import text_type, binary_type

def find_projection_vec(image_array, proj_dir='HORI', hori_th=0, vert_th=0):
    u"""
    :param image_array: 黑底白字
    :param proj_dir: 水平垂直投影的方向
    :return: 投影值列表
    """
    height, width = image_array.shape
    proj_vec = []
    if proj_dir is 'HORI':
        for id in range(height):
            pos = np.where(image_array[id, :] != 0)
            white_pixel_num = pos[0].shape[0]
            if white_pixel_num >= hori_th:
                proj_vec.append(white_pixel_num)
            else:
                proj_vec.append(0)

    elif proj_dir is 'VERT':
        for id in range(width):
            pos = np.where(image_array[:, id] != 0)
            white_pixel_num = pos[0].shape[0]
            if white_pixel_num >= vert_th:
                proj_vec.append(white_pixel_num)
            else:
                proj_vec.append(0)

    return proj_vec


def create_text_image(coordinate_list, height, width):
    result = np.zeros((height, width), np.uint8)
    for rect_id in range(len(coordinate_list)):
        xmin = coordinate_list[rect_id][0]
        ymin = coordinate_list[rect_id][1]
        xmax = coordinate_list[rect_id][2]
        ymax = coordinate_list[rect_id][3]
        result[ymin:ymax, xmin:xmax] = 255

    return result


def refinement_rect(image_text_rect_binary, image_text):
    u"""
    功能：
    1）对最原始的pixel link的结果进行边框精修，精修的过程利用了纯文本图像中的信息，
          因为文本图像中没有表格线和噪声。
    2）同时对同一文本行中间距很小的pixellink进行连接。
    3）连接标号问题

    参数：
    image_text_rect_binary: 最原始的pixel link检测结果 （黑底白字）
    image_text: 经过启发式过滤的纯文本图像（黑底白字）
    返回：
    image_refinement_rect_binary: 精修后的结果（白底黑字）
    """
    cc_info = find_connection_region(image_text_rect_binary)
    text_rect_list = cc_info['CC_POS']
    height, width = image_text_rect_binary.shape
    refinement_coordinate_list = []
    refinement_rect_area_list = []
    for id, rect in enumerate(text_rect_list):
        lx, ly, rx, ry = (
         rect[0], rect[1], rect[2], rect[3])
        image_crop = image_text[ly:ry, lx:rx]
        vertical_projection_vec = find_projection_vec(image_crop, 'VERT', hori_th=0, vert_th=2)
        horizontal_projection_vec = find_projection_vec(image_crop, 'HORI', hori_th=2, vert_th=0)
        vert_th, hori_th = (2, 2)
        v_candidate_pt_lst = np.where(np.array(vertical_projection_vec) != 0)[0]
        if len(v_candidate_pt_lst) == 0:
            continue
        v_seg_start = min(v_candidate_pt_lst) + lx
        v_seg_end = max(v_candidate_pt_lst) + lx
        h_candidate_pt_lst = np.where(np.array(horizontal_projection_vec) != 0)[0]
        if len(h_candidate_pt_lst) == 0:
            continue
        h_seg_start = min(h_candidate_pt_lst) + ly
        h_seg_end = max(h_candidate_pt_lst) + ly
        refinement_coordinate_list.append([v_seg_start, h_seg_start, v_seg_end, h_seg_end])
        rect_area = (h_seg_end - h_seg_start) * (v_seg_end - v_seg_start)
        refinement_rect_area_list.append(rect_area)

    if len(refinement_rect_area_list) != 0:
        max_rect_area = max(refinement_rect_area_list)
        for ii, val in enumerate(refinement_coordinate_list[:]):
            if refinement_rect_area_list[ii] * 200 <= max_rect_area:
                refinement_coordinate_list.remove(val)

    image_refinement_rect_binary = create_text_image(refinement_coordinate_list, height, width)
    image_refinement_rect_binary = 255 - image_refinement_rect_binary
    return (
     image_refinement_rect_binary, refinement_coordinate_list)


def find_connection_region(image_binary):
    u"""
    输入：
    image_binary: 希望提取连通区域的二值化图像（黑底白字）
    提取连通区域的相关函数：
    label(...), regionprops(...)
    http://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label
    http://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.regionprops
    返回：
    cc_info：连通区域的相关信息
    """
    cc_rect_pos = []
    cc_area = []
    cc_rect_area = []
    cc_density = []
    cc_ratio = []
    cc_coordinate = []
    label_image = label(image_binary, neighbors=8, connectivity=2)
    all_connect_info = regionprops(label_image)
    for element in all_connect_info:
        ly, lx, ry, rx = element.bbox
        lx = int(lx)
        ly = int(ly)
        rx = int(rx)
        ry = int(ry)
        rw = rx - lx
        rh = ry - ly
        short_side = float(min(rx - lx, ry - ly))
        long_side = float(max(rx - lx, ry - ly))
        area = element.area
        cc_rect_pos.append([lx, ly, rx, ry])
        cc_area.append(area)
        cc_rect_area.append(rw * rh)
        cc_density.append(float(area) / float(rw * rh))
        cc_ratio.append(short_side / long_side)
        cc_coordinate.append(element.coords)

    cc_info = {'CC_POS': cc_rect_pos, 'CC_AREA': cc_area, 'CC_RECT_AREA': cc_rect_area, 'CC_DENSITY': cc_density, 'CC_RATIO': cc_ratio, 'CC_COORDINATE': cc_coordinate}
    return cc_info


def image_resize(image, max_edge=1280):
    max_size = image.shape[0] if image.shape[0] > image.shape[1] else image.shape[1]
    if max_size > max_edge:
        resize_ratio = max_edge / float(max_size)
        return cv2.resize(image, None, fx=resize_ratio, fy=resize_ratio)
    else:
        return image
        return


def pil_rgba2rgb(image):
    bg = Image.new('RGB', image.size, (255, 255, 255))
    fg = image.split()[3]
    bg.paste(image, mask=fg)
    return bg


def is_np_image(image):
    return 'numpy' in str(type(image))


def is_pil_image(image):
    u"""是否是PIL格式的图像"""
    return 'PIL' in str(type(image))


def bgra2bgr(cv_img, background_color=(255, 255, 255)):
    u"""
    四通道图像转三通道图像
    :param cv_img:
    :param background_color:
    :return:
    """
    if len(cv_img.shape) < 3 and cv_img.shape[2] != 4:
        raise ValueError('The input should be an image with 4 channel')
    cv_img = cv_img.astype(np.float32)
    cv_img /= 255
    h, w, c = cv_img.shape
    background_color = list(background_color)
    background_color = np.array(background_color, dtype=np.float32) / 255
    background_img = np.ones((h, w, 3), dtype=np.float32) * background_color
    alpha = cv_img[:, :, 3]
    alpha = np.dstack([alpha, alpha, alpha])
    overlay = cv_img[:, :, :3] * alpha + background_img * (1 - alpha)
    overlay *= 255
    overlay = overlay.astype(np.uint8)
    return overlay


def np_image_to_pil(image):
    if len(image.shape) == 2:
        rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 3:
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    elif image.shape[2] == 4:
        rgb = bgra2bgr(image)
        rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
    else:
        raise Exception('Unsupported image channel: %d' % image.shape[2])
    return Image.fromarray(rgb)


def image_to_pil(image):
    if is_np_image(image):
        return np_image_to_pil(image)
    else:
        return image


def image_to_np(image):
    if is_pil_image(image):
        return pil_image_to_np(image)
    else:
        return image


def pil_image_to_np(image):
    if 'A' in image.getbands():
        image = pil_rgba2rgb(image)
    np_img = np.array(image.convert('RGB'))
    return cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)


def img_get_binary_data(image):
    image = image_to_pil(image)
    data = io.BytesIO()
    image.save(data, format='PNG')
    return data.getvalue()


def get_image_and_data(image):
    u"""
    获取numpy图像
    :param image:
    :return:
    """
    if isinstance(image, text_type):
        image = Image.open(image)
        if 'A' in image.getbands():
            image = pil_rgba2rgb(image)
        image = image_to_np(image)
    elif isinstance(image, binary_type):
        image = Image.open(io.BytesIO(image))
        if 'A' in image.getbands():
            image = pil_rgba2rgb(image)
        image = image_to_np(image)
    elif isinstance(image, np.ndarray):
        if image.shape[2] == 4:
            image = bgra2bgr(image)
    elif isinstance(image, Image.Image):
        image = image_to_np(image)
    else:
        raise TypeError(('image with type {} not support').format(type(image)))
    return image


def pil_rgba2rgb(image):
    bg = Image.new('RGB', image.size, (255, 255, 255))
    fg = image.split()[3]
    bg.paste(image, mask=fg)
    return bg


def get_image(imagefile):
    if isinstance(imagefile, text_type):
        image = Image.open(imagefile)
    elif isinstance(imagefile, binary_type):
        image = Image.open(io.BytesIO(imagefile))
    elif isinstance(imagefile, np.ndarray):
        image = Image.fromarray(imagefile)
    elif isinstance(imagefile, Image.Image):
        image = imagefile
    else:
        image = Image.new('RGB', (224, 224), 'black')
    if image.mode == 'RGBA':
        image = pil_rgba2rgb(image)
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    image = image.convert('RGB')
    return image


def save_images(images, save_dir):
    u"""
    保存图片集
    :param images:
    :param save_dir:
    :return:
    """
    for i, image in enumerate(images):
        shutil.copy(image, save_dir)


def is_too_small(image):
    w, h = image.size
    if w < 10 or h < 10 or w < 100 and h < 100:
        return True
    return False


def cv_to_pil(data, format='JPEG'):
    if format == 'PNG':
        img_encode = cv2.imencode('.png', data)
    else:
        img_encode = cv2.imencode('.jpg', data)
    return img_encode[1].tostring()


def bytes2cv(data):
    if isinstance(data, text_type):
        img_decode = cv2.imread(data)
    else:
        nparr = np.fromstring(data, np.uint8)
        img_decode = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    return img_decode


def list_images(obj_dir):
    IMAGE_SUFFIX = [
     '.jpg', '.JPG', '.jpeg', '.png', '.tiff', '.bmp', '.tif', '.TIF']
    paths = []
    for file_and_dir in os.listdir(obj_dir):
        if os.path.splitext(file_and_dir)[1] in IMAGE_SUFFIX:
            paths.append(os.path.abspath(os.path.join(obj_dir, file_and_dir)))

    return paths